'''
Malthus模型和Logistic模型
例题：
通过近两个世纪美国人口统计数据，建立模型，预测2010年人口数
模型建立:
设x(t)为第t年人口数量
设r(x)为人口年增长率 r(x)=r-sx
设xm为最大人口数，当x=xm时，r(xm)=0
得到Logistic人口模型：
    dx/dt = r(1-x/xm)x
    x(t0)=x0
其解为
    x(t)=xm/(1+((xm/x0)-1)e^(-r(t-t0)))
'''
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import leastsq

# 非线性最小二乘估计
data = pd.read_csv('p_data.csv')
x = data['year']
y = data['num']
t0 = x[0]
x0 = y[0]
Xi=np.array(x)
Yi=np.array(y)
fun=lambda p,x,y:(p[0] / (1 + ((p[0] / x0) - 1) * np.exp(-p[1] * (x - t0))))-y
# 参数的初始值会影响拟合效果？？？！！
p0 = np.array([300, 0.0274])
res = leastsq(fun, x0=p0, args=(Xi,Yi))
print(res)
print(fun(res[0],Xi,Yi))
#画样本点
plt.figure(figsize=(8,6)) ##指定图像比例： 8：6
plt.scatter(Xi,Yi,color="green",label="sample point",linewidth=2)
print(fun(res[0],2020,0))
#画拟合直线
x=np.linspace(1790,2020,1000) ##在0-15直接画100个连续点
y=fun(res[0],x,0)
print(x)
print(y)
plt.title('Population model')
plt.plot(x,y,color="red",label="fitting line",linewidth=2)
plt.legend(loc='lower right') #绘制图例
plt.show()

# 线性最小二乘估计
# 差分方程拟合
det=10
yp=y[1:21]
yy=y[0:20]
func=lambda p,yy,yp:((yp-yy)/det)*(1/yy)-p[0]-p[1]*yy
resc=leastsq(func,p0,args=(yy,yp))
print(res[0])
